#! /bin/env python

BLOCK_HEAD = '112'

KUAIJOB_EXE = r'D:\Madlee\Job\Madlee\kuai\bin32w\kuaijob.exe'
DATA_FOLDER = [
r'D:\Madlee\Job\CPMD_Data\FOX-7\2500K', 
r'D:\Madlee\Job\CPMD_Data\FOX-7\3000K',
r'D:\Madlee\Job\CPMD_Data\FOX-7\3500K',
r'D:\Madlee\Job\CPMD_Data\FOX-7\4000K']

MOLTYPE_FOLDER = r'D:\Madlee\Job\CPMD_Data\FOX-7\repository.mols'
REPOSITORY_DB = r'D:\Madlee\Job\CPMD_Data\FOX-7\repository.db'

XYZ_FILE = 'TRAJEC.xyz'
FRAME_GAP = 10


TRAJEC_DB = 'trajec.db'
PBC = 'PBC:  13.8680    6.6228   11.3119   90.0000   90.0650   90.0000\n'

import os
import os.path as path
from sqlite3 import connect
from subprocess import Popen, PIPE

KUAIJOB_SCRIPT = """job=parse_trajec
    database=%(trajec_db)s
    repository=%(repository_db)s
    molfolder=%(moltypes_folder)s
    synchronous=false
    cache_size=200000
    journal_mode=off
endjob
job=trace_mollife
    database=%(trajec_db)s
    repository=%(repository_db)s
    synchronous=false
    cache_size=200000
    journal_mode=off
    na = 0
endjob
job=trace_reaction
    database=%(trajec_db)s
    repository=%(repository_db)s
    synchronous=false
    cache_size=200000
    journal_mode=off
endjob
"""
SQL_DUMP_TXT = [("mol_info.txt", 
"""SELECT molecules.id as id, natoms, nbonds, formula, weight, moltype, 
    moltypes.data AS inchi_code, min(frame) as birth, max(frame) as death, 
    count(frame)*10 as life 
    FROM molecules INNER JOIN repository.moltypes 
	ON molecules.moltype = moltypes.id INNER JOIN mol_in_frame 
	ON molecules.id = mol_in_frame.molecule 
    GROUP BY (molecules.id);"""),
    ("mol_type.txt",
"""SELECT moltype, count(*) as n, natoms, nbonds, formula,  weight, inchi_code, 
    min(birth) as first_birth, avg(birth) as avg_birth, 
    max(death) as last_death, avg(death) as avg_death,
    min(life) as min_life, max(life) as max_life, avg(life) as avg_life, sum(life) as total_life
    FROM (SELECT molecules.id as id, natoms, nbonds, formula, weight, moltype, 
        moltypes.data AS inchi_code, min(frame) as birth, 
	max(frame) as death, count(frame)*10 as life 
	FROM molecules INNER JOIN repository.moltypes 
	    ON molecules.moltype = moltypes.id INNER JOIN mol_in_frame 
	    ON molecules.id = mol_in_frame.molecule 
	GROUP BY (molecules.id)
    ) GROUP by moltype;"""),
    ("frame_statics.txt", 
"""SELECT frame, moltype, formula, count(*) as hits
    FROM molecules INNER JOIN repository.moltypes 
	ON molecules.moltype = moltypes.id INNER JOIN mol_in_frame 
	ON molecules.id = mol_in_frame.molecule 
    GROUP BY frame, moltype ORDER BY frame, moltype;"""),
("rxn_types.txt", 
"""SELECT rxntypes.id as id, formula, data, count(*) as hits 
    FROM reactions INNER JOIN rxntypes 
        ON reactions.rxntype = rxntypes.id 
    GROUP BY rxntype ORDER BY rxntype;""",),
("mol_life.txt", """SELECT * from mol_life;"""),
]

SQL_SET_CREATE_FRAMES_TABLE = ["""CREATE TABLE frames (
    frame INTEGER PRIMARY KEY,
    data TEXT
);
"""]

SQL_INSERT_FRAME = """INSERT INTO frames (frame, data) VALUES (?, ?)"""

def create_db(filename):
    if path.isfile(filename):
        raise RuntimeError("%s is existed in current folder." % filename)

    conn = connect(filename)
    cursor = conn.cursor()

    data = globals()
    for i in SQL_SET_CREATE_FRAMES_TABLE:
        cursor.execute(i % data)
    conn.commit()
    del cursor
    return conn

def save_frame(cursor, frame, data):
    cursor.execute(SQL_INSERT_FRAME, (frame, data))

def add_frame(cursor, block):
    assert len(block) == int(block[0].strip()) + 2
    frame = int(block[1].strip().split()[1])
    block[1] = block[1].rstrip() + PBC
    save_frame(cursor, frame, ''.join(block))
    print frame

def clean_up():
    if path.isfile(REPOSITORY_DB):
        os.remove(REPOSITORY_DB)
    if path.isdir(MOLTYPE_FOLDER):
        for i in os.listdir(MOLTYPE_FOLDER):
            if i.endswith('.mol'):
                os.remove(path.join(MOLTYPE_FOLDER, i))
    else:
        os.makedirs(MOLTYPE_FOLDER)

    for i in DATA_FOLDER:
        filename = path.join(i, TRAJEC_DB)
        if path.isfile(filename):
            os.remove(filename)
        

def import_traj(folder):
    dataname = path.join(folder, XYZ_FILE)
    dbname = path.join(folder, TRAJEC_DB)
    print "Importing", dataname, "into", dbname
    db = create_db(dbname)
    cursor = db.cursor()
    block = None
    with open(dataname) as file:
        for i in file:
            if i.strip() == BLOCK_HEAD:
                if block:
                    add_frame(cursor, block)
                block=[i]
            else:
                block.append(i)
    add_frame(cursor, block)
    db.commit()
    
    print "Job Done"

def parse_traj_db(folder):
    trajec_db = path.join(folder, TRAJEC_DB)
    print "Parsing", trajec_db
    repository_db = REPOSITORY_DB
    moltypes_folder = MOLTYPE_FOLDER
    scripts = KUAIJOB_SCRIPT % locals()
    print scripts
    pipe = Popen("kuaijob", executable=KUAIJOB_EXE, stdin=PIPE)
    pipe.stdin.write(scripts)
    pipe.stdin.close()
    pipe.wait()


def dump_info(folder, sep='|'):
    dbname = path.join(folder, TRAJEC_DB)
    
    conn = connect(dbname)
    cursor = conn.cursor()
    cursor.execute("ATTACH DATABASE ? as repository;", [REPOSITORY_DB])
    for i, j in SQL_DUMP_TXT:
        cursor.execute(j)
        txtfile = path.join(folder, i)
        with open(txtfile, "wt") as file:
            print "Dumping ", txtfile
            title = sep.join([i[0] for i in cursor.description])
            file.write(title+"\n")
            for row in cursor:
                file.write(sep.join([str(k) for k in row])+"\n")

def hits(file):
    resultR = {}
    resultP = {}
    resultR[1] = resultP[1] = 0
    file.next()
    for i in file:
        tokens = i.split('|')
        n = int(tokens[3].strip())
        r, p = tokens[2].split('=')
        
        for j in r.split('+'):
            j = j.strip()[:-1]
            if j[0] == '[':
                factor = 1
                j = j[1:]
            else:
                factor = int(j[0:j.find('[')])
                j = j[j.find('[')+1:]
            key = int(j)
            if key in resultR:
                resultR[key] += n*factor
            else:
                resultR[key] = n*factor
                
        for j in p.split('+'):
            j = j.strip()[:-1]
            if j[0] == '[':
                factor = 1
                j = j[1:]
            else:
                factor = int(j[:j.find('[')])
                j = j[j.find('[')+1:]
            key = int(j)
            if key in resultP:
                resultP[key] += n*factor
            else:
                resultP[key] = n*factor

    return resultR, resultP

def modify_hits(folder, sep="|"):
    with open(path.join(folder, "rxn_types.txt")) as file:
        r, p = hits(file)
    lines = []
    with open(path.join(folder, "mol_type.txt")) as file:
        for line in file:
            lines.append(line.rstrip())
    
    with open(path.join(folder, "mol_type.txt"), "wt") as file:
        file.write(sep.join([lines[0].rstrip(), 'r','p','-','activity']) + "\n")
        for line in lines[1:]:
            id, tokens = line.split(sep, 1)
            id = int(id)
            if id not in r:
                r[id] = 0
            if id not in p:
                p[id] = 0
            assert r[id] + p[id] > 0
            file.write(sep.join(str(i) for i in [id, tokens, r[id], p[id], r[id]-p[id], r[id]+p[id]])+"\n")
            

def clean_rxn(folder, sep='|'):
    print "Cleaning", path.join(folder, "rxn_types.txt")
    line2 = None
    fileI = open(path.join(folder, "rxn_types.txt"))
    fileO = open(path.join(folder, "clean_rxn.txt"), "wt")
    fileO.write(fileI.next())
    try:
        line1 = fileI.next().split(sep)
        while True:
            line2 = fileI.next().split(sep)
            if int(line1[0]) + 1 == int(line2[0]):
                n1 = int(line1[3])
                n2 = int(line2[3])
                if n1 > n2:
                    line1[3] = str(n1 - n2)
                    fileO.write(sep.join(line1)+'\n')
                elif n1 < n2: 
                    line2[3] = str(n2 - n1)
                    fileO.write(sep.join(line2)+'\n')
                line1 = fileI.next().split(sep)
                line2 = None
            else:
                fileO.write(sep.join(line1))
                line1 = line2
    except StopIteration:
        pass
        
    if line2:
        fileO.write(sep.join(line2))
        
if __name__ == '__main__':
    clean_up()
    for i in DATA_FOLDER:
        print "    **************    Parsing %s     **************" % i
        import_traj(i)
        parse_traj_db(i)
        dump_info(i)
        modify_hits(i)
        clean_rxn(i)
        
